import matplotlib.pyplot as plt
import numpy as np

from rootPath import project_path
from utils.drawUtils import paresJsonObject

# 生成示例数据
x = range(1, 201)
file_path = "../mr/logs/sdss/2024_10_04_08_25_20/time_records.json"
file_path1 = "../mr_dd/logs/sdss/2024_10_04_20_35_54/time_records.json"
file_path2 = "../mr_dd_softmax/logs/sdss/2024_10_05_06_39_15/time_records.json"
ignore_path1 = "../mr_dd/logs/sdss/2024_10_04_20_35_54/ignore_rate_record.json"
ignore_path2 = "../mr_dd_softmax/logs/sdss/2024_10_05_06_39_15/ignore_rate_record.json"
json_obj = paresJsonObject(file_path)
json_obj1 = paresJsonObject(file_path1)
json_obj2 = paresJsonObject(file_path2)
ignore_obj = paresJsonObject(ignore_path1)
ignore_obj2 = paresJsonObject(ignore_path2)
epoch_com_acc_list = json_obj["epoch_com_acc_list"]
epoch_com_acc_list1 = json_obj1["epoch_com_acc_list"]
epoch_com_acc_list2 = json_obj2["epoch_com_acc_list"]
ignore_com_acc_list1 = ignore_obj["epoch_com_acc_list"]
ignore_com_acc_list2 = ignore_obj2["epoch_com_acc_list"]
mr_times = []
mr_dd_times = []
mr_dd_softmax_times = []
ignore_dd_percentage = []
ignore_dd_percentage2 = []
for i in range(1, len(epoch_com_acc_list)):
    mr_times.append(epoch_com_acc_list[i]["(2,)"])
    mr_dd_times.append(epoch_com_acc_list1[i]["(2,)"])
    mr_dd_softmax_times.append(epoch_com_acc_list2[i]["(2,)"])
    temp = 0.0
    temp2 = 0.0
    for key in ignore_com_acc_list1[i].keys():
        temp += ignore_com_acc_list1[i][key]
        temp2 += ignore_com_acc_list2[i][key]
    temp = temp / len(ignore_com_acc_list1[i].keys())
    temp2 = temp2 / len(ignore_com_acc_list1[i].keys())
    ignore_dd_percentage.append(temp)
    ignore_dd_percentage2.append(temp2)


# print(sum(mr_times))
# print(sum(mr_dd_times))
# print(sum(mr_dd_softmax_times))

# 创建图形和主坐标轴
fig, ax1 = plt.subplots(figsize=(10, 6))

# 绘制第一条线（使用主坐标轴）
colorB = 'tab:black'
color1 = 'tab:blue'
color2 = 'tab:green'  # 为第二条线选择不同的颜色
color3 = 'tab:orange'  # 为第二条线选择不同的颜色
ax1.set_xlabel('epoch',fontsize=15)
ax1.set_ylabel('Contribution Measurement Time', fontsize=15)
ax1.plot(x, mr_times, color=color1, label='MR Times')
ax1.plot(x, mr_dd_times, color=color2, label='MR RS Times')
ax1.plot(x, mr_dd_softmax_times, color=color3, label='MR WS Times')
# ax1.tick_params(axis='y')
ax1.tick_params(axis='both', labelsize=15)  # 主坐标轴
ax1.set_ylim(58, 80)  # 设置y轴范围
# 添加图例
ax1.legend()
# 创建次坐标轴
ax2 = ax1.twinx()

# 绘制第二条线（使用次坐标轴）
color = 'tab:orange'
ax2.set_ylabel('Percentage of pruning',fontsize=15)
ax2.plot(x, ignore_dd_percentage, linestyle='--', color=color2, label='RS')
ax2.plot(x, ignore_dd_percentage2, linestyle='--', color=color3, label='WS')
ax2.tick_params(axis='both', labelsize=15)  # 主坐标轴
ax2.set_ylim(0, 76)  # 设置y轴范围
# 添加标题
plt.title('')

# 添加图例
lines1, labels1 = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax1.legend(lines1, ['MR', 'MR+RS', 'MR+WS'], loc='upper left')
ax2.legend(lines2, ['RS', "WS"], loc='upper right')
# 调整布局并显示图形
fig.tight_layout()
# 保存图形到文件，确保文件名包含 .png 扩展名
plt.savefig(project_path + '/experiment/imgs/time_eval_epoch.png', dpi=300, bbox_inches='tight')
plt.show()
